#include "stft.h"
#include <Arduino.h>

#define winType 0.5

#define KISSSIZE 256





kiss_fftr_cfg stftInit(){
    kiss_fftr_cfg kfft_cfg;
    void*scratch;
size_t scratch_size = 0;
  kfft_cfg = kiss_fftr_alloc(KISSSIZE, 0, nullptr, &scratch_size);
  if (kfft_cfg != nullptr) {
    Serial.println( "Kiss memory sizing failed.\n");
    return NULL;
  }
  scratch = malloc(scratch_size);
  if (scratch == nullptr) {
    Serial.println( "Failed to alloc fft scratch buffer\n");
    return NULL;
  }
  // Let kissfft configure the scratch space we just allocated
  kfft_cfg = kiss_fftr_alloc(KISSSIZE, 0,scratch, &scratch_size);
  if(kfft_cfg==nullptr){
    Serial.println( "config failed\n");
    return NULL;
  }
    Serial.printf("scratch_size:%d\r\n",scratch_size);
    return kfft_cfg;
}

void hammingWindow(float*out,int window_size){
	const float arg = M_PI * 2.0;
	for(int i=0;i<window_size;i++){
		float float_value = winType - ((1-winType)* cos(arg*i/window_size ));
		out[i]=out[i]*float_value;
	}
}

void stft(const uint16_t*left_data,float*out,kiss_fftr_cfg kfft_cfg){
     //创建一个等长的float数组，将int16转换成float
     float data[KISSSIZE]={0};
     kiss_fft_cpx outData[KISSSIZE/2+1]={0,0};
     for(int i=0;i<KISSSIZE;i++){
        int16_t temp=left_data[i];
        data[i]=temp/32768.0;
     }
     //加窗
     hammingWindow(data,KISSSIZE);
     //快速傅里叶变换
     kiss_fftr(kfft_cfg,data,outData);
     for(int i=0;i<KISSSIZE/2+1;i++){
         //out[i]=abs(outData[i].r);
        //这里不是简单的求绝对值，是虚部和实部平方然后相加
         float t=(outData[i].r*outData[i].r)+(outData[i].i*outData[i].i);
         out[i]=sqrt(t);
     }

}

void stftAll(const uint16_t*left_data,float*out,kiss_fftr_cfg kfft_cfg){
    for(int i=0;i<124;i++){
      //窗口每次移动128步
      const uint16_t*input=left_data+i*128;
      //结果集每次移动129个
      float*myout=out+i*129;
      stft(input,myout,kfft_cfg);
    }
}


